// **********************************************************
// >> File Name: /home/diabio/project/thread_pool/include/thread_pool.c
// >> Author: diabiolin
// >> Created Time: Tue 29 Apr 2025 09:46:46 AM CST
// >> Description:
// **********************************************************

#define _POSIX_C_SOURCE 199309L
#include "thread_pool.h"
#include <stdlib.h>
#include <errno.h>
#include <time.h>
#include <string.h>

// 函数声明
static void threadPoolFree(ThreadPool_t* threadpool);
static int threadPoolStart(ThreadPool_t* threadPool, int createNum);

// 线程回调函数（等待线程池分配任务并执行分配的任务）
void* threadCallBack(void* arg){
    ThreadPool_t* threadPool = (ThreadPool_t*)arg;
    ThreadTask_t task;

    while(1){
        pthread_mutex_lock(&threadPool->mutex);
        threadPool->threadWaitNum++;  // 进入等待状态
        
        while(threadPool->taskWaitNum == 0 && !threadPool->isExit){
            pthread_cond_wait(&threadPool->cond, &threadPool->mutex);
        }

        // 检查是否需要退出
        if(threadPool->isExit){
            if(!(threadPool->shutdownMode == THREAD_POOL_GRACEFUL_SHUTDOWN && 
                 threadPool->taskWaitNum > 0)){
                threadPool->threadWaitNum--;
                pthread_mutex_unlock(&threadPool->mutex);
                break;
            }
        }

        // 获取任务前减少等待线程数
        threadPool->threadWaitNum--;

        // 从任务队列中取出任务
        task.func = threadPool->taskQueue[threadPool->taskQueueHead].func;
        task.args = threadPool->taskQueue[threadPool->taskQueueHead].args;

        threadPool->taskQueueHead = (threadPool->taskQueueHead + 1) % threadPool->taskQueueSize;
        threadPool->taskWaitNum--;
        threadPool->taskRunningNum++;
        pthread_mutex_unlock(&threadPool->mutex);

        // 执行任务
        if(task.func != NULL){
            task.func(task.args);
        }

        // 任务完成后更新计数
        pthread_mutex_lock(&threadPool->mutex);
        threadPool->taskRunningNum--;

        // 如果是优雅关闭模式且所有任务都完成了，发出信号
        if(threadPool->shutdownMode == THREAD_POOL_GRACEFUL_SHUTDOWN && 
           threadPool->taskWaitNum == 0 && 
           threadPool->taskRunningNum == 0){
            pthread_cond_broadcast(&threadPool->cond);
        }
        pthread_mutex_unlock(&threadPool->mutex);
    }

    // 线程退出前的清理工作
    pthread_mutex_lock(&threadPool->mutex);
    threadPool->threadStartNum--;
    pthread_mutex_unlock(&threadPool->mutex);
    return NULL;
}

// 创建线程
ThreadPool_t* threadPoolCreate(int threadMaxNum, int threadStartNum, int taskQueueSize){
    if(threadMaxNum <= 0 || threadStartNum <= 0 || taskQueueSize <= 0 || threadStartNum > threadMaxNum){
        return NULL;
    }

    ThreadPool_t* threadPool = (ThreadPool_t*)calloc(1, sizeof(ThreadPool_t));
    if(threadPool == NULL){
        return NULL;
    }

    // 初始化基本参数
    threadPool->threadMaxNum = threadMaxNum;
    threadPool->threadStartNum = 0;
    threadPool->taskQueueSize = taskQueueSize;
    threadPool->isExit = 0;
    threadPool->shutdownMode = 0;
    threadPool->taskRunningNum = 0;
    threadPool->taskWaitNum = 0;
    threadPool->taskQueueHead = 0;
    threadPool->taskQueueTail = 0;
    threadPool->threadWaitNum = 0;

    // 分配线程数组内存
    threadPool->threads = (pthread_t*)calloc(threadMaxNum, sizeof(pthread_t));
    if(threadPool->threads == NULL){
        threadPoolFree(threadPool);
        return NULL;
    }

    // 分配任务队列内存
    threadPool->taskQueue = (ThreadTask_t*)calloc(taskQueueSize, sizeof(ThreadTask_t));
    if(threadPool->taskQueue == NULL){
        threadPoolFree(threadPool);
        return NULL;
    }

    // 初始化互斥锁和条件变量
    if(pthread_mutex_init(&threadPool->mutex, NULL) != 0){
        threadPoolFree(threadPool);
        return NULL;
    }

    if(pthread_cond_init(&threadPool->cond, NULL) != 0){
        pthread_mutex_destroy(&threadPool->mutex);
        threadPoolFree(threadPool);
        return NULL;
    }

    // 创建初始线程
    if(threadPoolStart(threadPool, threadStartNum) <= 0){
        threadPoolDestroy(threadPool);
        return NULL;
    }

    return threadPool;
}

// 添加任务
int threadPoolAddTask(ThreadPool_t* threadPool, TASKFUNC taskfunc, void* args){
    if(threadPool == NULL || taskfunc == NULL){
        return THREAD_POOL_INVALID_ARG;
    }

    pthread_mutex_lock(&threadPool->mutex);

    // 检查线程池是否正在关闭
    if(threadPool->isExit){
        pthread_mutex_unlock(&threadPool->mutex);
        return THREAD_POOL_SHUTDOWN;
    }

    // 检查任务队列是否已满
    if((threadPool->taskQueueTail + 1) % threadPool->taskQueueSize == threadPool->taskQueueHead){
        // 如果当前线程数小于最大线程数，尝试创建新线程
        if(threadPool->threadStartNum < threadPool->threadMaxNum){
			// 先解锁，再创建新线程（这个操作自己会加锁），再加锁，防止死锁
            pthread_mutex_unlock(&threadPool->mutex);
            threadPoolStart(threadPool, 1);
            pthread_mutex_lock(&threadPool->mutex);
        }
        pthread_mutex_unlock(&threadPool->mutex);
        return THREAD_POOL_QUEUE_FULL;
    }

    // 添加任务到队列
    threadPool->taskQueue[threadPool->taskQueueTail].func = taskfunc;
    threadPool->taskQueue[threadPool->taskQueueTail].args = args;
    threadPool->taskQueueTail = (threadPool->taskQueueTail + 1) % threadPool->taskQueueSize;
    threadPool->taskWaitNum++;

    // 唤醒一个等待的线程
    pthread_cond_signal(&threadPool->cond);

    pthread_mutex_unlock(&threadPool->mutex);
    return THREAD_POOL_SUCCESS;
}

// 创建线程
static int threadPoolStart(ThreadPool_t* threadPool, int createNum){
    if(threadPool == NULL || createNum <= 0){
        return THREAD_POOL_INVALID_ARG;
    }

    pthread_mutex_lock(&threadPool->mutex);
    
    int curThreadNum = threadPool->threadStartNum;
    int targetThreadNum = curThreadNum + createNum;
    if(targetThreadNum > threadPool->threadMaxNum){
        targetThreadNum = threadPool->threadMaxNum;
    }

    int successCount = 0;
    for(int i = curThreadNum; i < targetThreadNum; i++){
        if(pthread_create(&threadPool->threads[i], NULL, threadCallBack, threadPool) != 0){
            break;
        }
        successCount++;
        threadPool->threadStartNum++;
    }

    pthread_mutex_unlock(&threadPool->mutex);
    return successCount > 0 ? successCount : THREAD_POOL_INVALID_ARG;
}

// 获取任务数量
int threadPoolGetTaskCount(ThreadPool_t* threadPool, int* waiting, int* running){
    if(threadPool == NULL || waiting == NULL || running == NULL){
        return THREAD_POOL_INVALID_ARG;
    }

    pthread_mutex_lock(&threadPool->mutex);
    *waiting = threadPool->taskWaitNum;
    *running = threadPool->taskRunningNum;
    pthread_mutex_unlock(&threadPool->mutex);

    return THREAD_POOL_SUCCESS;
}

// 检查线程池是否已关闭
int threadPoolIsShutdown(ThreadPool_t* threadPool){
    if(threadPool == NULL){
        return 1;
    }

    pthread_mutex_lock(&threadPool->mutex);
    int isExit = threadPool->isExit;
    pthread_mutex_unlock(&threadPool->mutex);

    return isExit;
}

// 立即关闭线程池
int threadPoolDestroy(ThreadPool_t* threadPool){
    if(threadPool == NULL){
        return THREAD_POOL_INVALID_ARG;
    }

    pthread_mutex_lock(&threadPool->mutex);
    if(threadPool->isExit){
        pthread_mutex_unlock(&threadPool->mutex);
        return THREAD_POOL_SHUTDOWN;
    }

    threadPool->shutdownMode = THREAD_POOL_IMMEDIATE_SHUTDOWN;
    threadPool->isExit = 1;
	int threads_to_join = threadPool->threadStartNum; // 捕获需要join的线程数
    pthread_mutex_unlock(&threadPool->mutex);

    // 唤醒所有等待的线程
    pthread_cond_broadcast(&threadPool->cond);

    // 等待所有线程退出
    for(int i = 0; i < threads_to_join; i++){
        pthread_join(threadPool->threads[i], NULL);
    }

    threadPoolFree(threadPool);
    return THREAD_POOL_SUCCESS;
}

// 优雅关闭线程池
int threadPoolGracefulDestroy(ThreadPool_t* threadPool){
    if(threadPool == NULL){
        return THREAD_POOL_INVALID_ARG;
    }

    pthread_mutex_lock(&threadPool->mutex);
    if(threadPool->isExit){
        pthread_mutex_unlock(&threadPool->mutex);
        return THREAD_POOL_SHUTDOWN;
    }

    threadPool->shutdownMode = THREAD_POOL_GRACEFUL_SHUTDOWN;
    threadPool->isExit = 1;

    // 等待所有任务完成，设置60秒超时
    struct timespec ts;
    clock_gettime(CLOCK_REALTIME, &ts);
    ts.tv_sec += 60;

    int ret = THREAD_POOL_SUCCESS;
    while(threadPool->taskWaitNum > 0 || threadPool->taskRunningNum > 0){
        int wait_ret = pthread_cond_timedwait(&threadPool->cond, &threadPool->mutex, &ts);
        if(wait_ret == ETIMEDOUT){
            ret = THREAD_POOL_TIMEOUT;
            break;
        }
    }
	int threads_to_join = threadPool->threadStartNum; // 捕获需要join的线程数
    pthread_mutex_unlock(&threadPool->mutex);

    // 唤醒所有等待的线程
    pthread_cond_broadcast(&threadPool->cond);

    // 等待所有线程退出
    for(int i = 0; i < threads_to_join; i++){
        pthread_join(threadPool->threads[i], NULL);
    }

    threadPoolFree(threadPool);
    return ret;
}

// 释放线程池资源
static void threadPoolFree(ThreadPool_t* threadpool){
    if(threadpool == NULL){
        return;
    }

    if(threadpool->threads != NULL){
        free(threadpool->threads);
        threadpool->threads = NULL;
    }

    if(threadpool->taskQueue != NULL){
        free(threadpool->taskQueue);
        threadpool->taskQueue = NULL;
    }

    pthread_mutex_destroy(&threadpool->mutex);
    pthread_cond_destroy(&threadpool->cond);

    free(threadpool);
}
